#!/usr/bin/env python

""" Example for detect dns denial of service attacks """

__author__ = "Luis Campo Giralte"
__copyright__ = "Copyright (C) 2020-2020 by Luis Campo Giralte"
__revision__ = "$Id$"
__version__ = "0.1"

import operator 
import sys
import pyaiengine

stack = None

def wait_handler():
    """Temporary callback for wait 30 seconds
    to restart check metrics again."""

    pd.remove_timer(30)

    # Reset the DNS counters/metrics
    stack.reset_counters("dns")

    # Enable the sampling again
    pd.add_timer(scheduler_reflection_dns, 5)

def scheduler_reflection_dns():
    """Main calllback that will check 
    every 5 seconds the status of the countes."""

    print("DNS attack DoS Checker")
    counters = stack.get_counters("dns")

    queries = int(counters['queries'])
    resp = int(counters['responses'])

    if abs(queries - resp) > 10000:
        data = stack.get_cache_data("DNS", "name")
        attack_domain = max(data.items(), key=operator.itemgetter(1))[0]
        print("Potential DNS attack on domain '%s'" % attack_domain)
        pd.remove_timer(5) 
        pd.add_timer(wait_handler, 30)

if __name__ == '__main__':

    # Load an instance of a Network Stack Lan
    stack = pyaiengine.StackLan()

    # Allocate 1M flows
    stack.udp_flows = 1000000

    # Create a instace of a PacketDispatcher
    with pyaiengine.PacketDispatcher("lo0") as pd:
        pd.stack = stack
        # Sets a handler method that will be call
        # every 5 seconds for check the values
        pd.add_timer(scheduler_reflection_dns, 5)
        pd.run()

    sys.exit(0)
